package com.aiwiown.snackmq.broker.handler.console;

import com.aiwiown.snackmq.broker.auth.AuthorizationService;
import com.aiwiown.snackmq.cluster.raft.SnackMQStateMachine;
import com.aiwiown.snackmq.cluster.raft.manager.SnackMQRaftNodeManager;
import com.aiwiown.snackmq.common.auth.Action;
import com.aiwiown.snackmq.common.dto.console.ConsumerGroupDTO;
import com.aiwiown.snackmq.common.dto.console.PartitionLagDTO;
import com.aiwiown.snackmq.common.handler.MessageHandler;
import com.aiwiown.snackmq.common.handler.RequestContext;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageStatus;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.common.serialization.JsonSerializer;
import com.aiwiown.snackmq.common.serialization.SerializerFactory;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Slf4j
public class ListConsumerGroupsHandler implements MessageHandler {

    private final SnackMQStateMachine stateMachine;
    private final SnackMQRaftNodeManager dataRaftNodeManager;
    private final AuthorizationService authorizationService;
    private final JsonSerializer serializer = SerializerFactory.getJsonSerializer();

    public ListConsumerGroupsHandler(SnackMQStateMachine stateMachine, 
                                   SnackMQRaftNodeManager dataRaftNodeManager,
                                   AuthorizationService authorizationService) {
        this.stateMachine = stateMachine;
        this.dataRaftNodeManager = dataRaftNodeManager;
        this.authorizationService = authorizationService;
    }

    @Override
    public MessageType getMessageType() {
        return MessageType.LIST_CONSUMER_GROUPS_REQUEST;
    }

    @Override
    public void handle(RequestContext ctx, Message request) {
        try {
            authorizationService.authorize(request.getProperty("authToken"), Action.MANAGE, new AuthorizationService.Resource("CLUSTER", "*"));

            // 检查当前节点是否为数据Raft组的Leader
            if (!dataRaftNodeManager.isLeader()) {
                String leaderAddress = dataRaftNodeManager.getLeaderAddress();
                if (leaderAddress != null) {
                    // 如果不是 Leader，则返回重定向响应
                    log.debug("Not the data raft leader. Redirecting LIST_CONSUMER_GROUPS_REQUEST to leader at {}", leaderAddress);
                    ctx.writeResponse(Message.builder()
                            .type(MessageType.RESPONSE)
                            .messageId(request.getMessageId())
                            .status(MessageStatus.REDIRECT)
                            .property("redirectAddress", leaderAddress)
                            .build());
                    return;
                }
                throw new IllegalStateException("No data raft leader available to serve the request.");
            }

            // 【修复】不再使用本地 StorageService。
            // 从状态机获取全局一致的消费者位点和分区高水位线。
            Map<String, Map<String, Long>> allOffsets = stateMachine.getAllConsumerGroupOffsets();
            Map<String, Long> allHighWatermarks = stateMachine.getAllPartitionHighWatermarks();
            List<ConsumerGroupDTO> resultList = new ArrayList<>();

            for (Map.Entry<String, Map<String, Long>> groupEntry : allOffsets.entrySet()) {
                String groupName = groupEntry.getKey();
                long totalLagForGroup = 0;
                // 【修复】使用新的 PartitionLagDTO 来存储每个分区的详细延迟信息，以支持多主题消费组。
                List<PartitionLagDTO> partitionLags = new ArrayList<>();

                for (Map.Entry<String, Long> offsetEntry : groupEntry.getValue().entrySet()) {
                    String topicPartitionKey = offsetEntry.getKey(); // 格式为 "topic:partitionId"
                    long committedOffset = offsetEntry.getValue();

                    // 【修复】从全局一致的 highWatermarks Map 中获取数据。
                    // 如果分区还没有任何消息，高水位线为-1，与初始提交位点-1一致，lag为0。
                    long highWatermark = allHighWatermarks.getOrDefault(topicPartitionKey, -1L);

                    // 计算延迟。如果高水位线小于已提交位点（异常情况），则延迟为0。
                    long lag = (highWatermark >= committedOffset) ? (highWatermark - committedOffset) : 0;

                    totalLagForGroup += lag;
                    partitionLags.add(new PartitionLagDTO(topicPartitionKey, committedOffset, highWatermark, lag));
                }

                // 【修复】使用新的 ConsumerGroupDTO 构造函数，它能正确处理总延迟和详细分区延迟。
                resultList.add(new ConsumerGroupDTO(groupName, totalLagForGroup, partitionLags));
            }

            Message response = Message.builder()
                    .messageId(request.getMessageId())
                    .type(MessageType.RESPONSE)
                    .status(MessageStatus.SUCCESS)
                    .body(serializer.serialize(resultList))
                    .build();
            ctx.writeResponse(response);

        } catch (Exception e) {
            log.warn("Failed to list consumer groups: {}", e.getMessage(), e);
            Message errorResponse = Message.builder()
                    .messageId(request.getMessageId())
                    .type(MessageType.RESPONSE)
                    .status(MessageStatus.FAILED)
                    .errorMessage(e.getMessage())
                    .build();
            ctx.writeResponse(errorResponse);
        }
    }
}